
#include <boost/beast/core.hpp>
#include <boost/beast/http.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <iostream>

namespace beast = boost::beast;
namespace http = beast::http;
namespace websocket = beast::websocket;
namespace net = boost::asio;
using tcp = net::ip::tcp;

void handle_http(http::request<http::string_body> req, tcp::socket& socket) {
    http::response<http::string_body> res{http::status::ok, req.version()};
    res.set(http::field::server, "Dual-Protocol Server");
    res.body() = "HTTP Request Handled";
    res.prepare_payload();
    http::write(socket, res);
}

void run_server(unsigned short port) {
    try {
        net::io_context ioc;
        tcp::acceptor acceptor{ioc, {tcp::v4(), port}};
        
        while(true) {
            tcp::socket socket{ioc};
            acceptor.accept(socket);
            
            // 检测是否为WebSocket升级请求
            beast::flat_buffer buffer;
            http::request<http::string_body> req;
            http::read(socket, buffer, req);
            
            if(websocket::is_upgrade(req)) {
                websocket::stream<tcp::socket> ws{std::move(socket)};
                ws.accept(req); // 完成WebSocket握手
                
                ws.read(buffer); // 接收WebSocket消息
                ws.text(ws.got_text());
                ws.write(buffer.data()); // 回传消息
            } else {
                handle_http(req, socket); // 处理普通HTTP请求
                socket.shutdown(tcp::socket::shutdown_send);
            }
        }
    } catch(const std::exception& e) {
        std::cerr << "Error: " << e.what() << std::endl;
    }
}

int main() {
    run_server(8080);
}